#!/usr/bin/env python3
'''
该q值更新是在每次游玩一步的过程中更新q值
'''
import gymnasium as gym
import collections
from tensorboardX import SummaryWriter

ENV_NAME = "FrozenLake-v1"
GAMMA = 0.9
ALPHA = 0.2
TEST_EPISODES = 20


class Agent:
    def __init__(self):
        self.env = gym.make(ENV_NAME)
        self.state = self.env.reset()
        # 以当前状态，以及执行的动作为key，保存获取的 激励
        self.values = collections.defaultdict(float)

    def sample_env(self):
        action = self.env.action_space.sample()
        old_state = self.state
        new_state, reward, is_done, _ = self.env.step(action)
        self.state = self.env.reset() if is_done else new_state
        return (old_state, action, reward, new_state)

    def best_value_and_action(self, state):
        '''
        根据状态获取下一个状态能得到最高的激励的动作和对应的激励值


        '''

        best_value, best_action = None, None
        # 遍历动作空间
        for action in range(self.env.action_space.n):
            # 根据状态和动作获取可获取的激励（保存起来的）
            action_value = self.values[(state, action)]
            # 如果best_value为空或者有更高的激励值，则保存当前的动作和激励
            # 如果状态为空，那么是随机返回一个0
            if best_value is None or best_value < action_value:
                best_value = action_value
                best_action = action
        return best_value, best_action

    def value_update(self, s, a, r, next_s):
        '''
        根据当前的动作，更新值

        s: 当前的状态
        a: 执行的动作
        r： 获取的激励
        next_s： 下一个状态
        '''
        # 获取next_s装填所能达到的最大的激励
        best_v, _ = self.best_value_and_action(next_s)
        # 当前的激励 + GAMMA * 下一个状态所能获取的最大激励
        new_val = r + GAMMA * best_v
        # 获取当前状态以及执行动作所能获取到的激励值
        old_val = self.values[(s, a)]
        # 这一行代码根据Q - learning的更新公式更新了当前状态 - 动作对(s, a)
        # 的Q值。这里，ALPHA是学习率，用于控制学习过程中新信息和旧信息的权重。公式old_val * (
        #             1 - ALPHA) + new_val * ALPHA是一个加权平均，其中旧Q值的权重为(
        #     1 - ALPHA)，新计算出的期望累积回报的权重为ALPHA。这种更新方式使得Q值在迭代过程中逐渐收敛到最优Q值。

        # 在Q-learning算法中，更新当前状态-动作对的Q值时，考虑新旧Q值的权重是为了平衡学习过程中新信息和旧信息的影响。具体来说，这里的权重由学习率（ALPHA）来调节。学习率决定了在每次更新时新信息对Q值的影响程度，较小的学习率意味着新信息对Q值的影响较小，而较大的学习率则表示新信息对Q值的影响较大。
        #
        # 当然，你提到的直接用新Q值覆盖旧Q值的方法也是可行的，即将学习率设置为1。但是这种方法存在一些缺点：
        #
        # 收敛速度：直接用新Q值覆盖旧Q值可能导致收敛速度较慢。因为这种更新方式相当于完全忽略了之前学到的信息，每次更新都会受到最近一次采样的较大影响。这可能会导致Q值在不同的估计之间波动较大，使得算法需要更多的迭代次数才能收敛到最优值。
        #
        # 鲁棒性：完全使用新Q值覆盖旧Q值可能导致算法的鲁棒性较差。在某些情况下，由于环境噪声、观测误差或其他因素，可能会得到一个不准确的新Q值。如果直接用这个不准确的新Q值覆盖旧Q值，可能会导致算法受到这些误差的影响，从而降低最终策略的性能。
        #
        # 综上所述，虽然直接用新Q值覆盖旧Q值是可行的，但在实际应用中，考虑新旧Q值的权重更为常见，因为这样可以平衡新信息和旧信息的影响，提高算法的收敛速度和鲁棒性。这也是为什么在Q-learning算法中，通常会采用加权平均的方式来更新Q值
        self.values[(s, a)] = old_val * (1-ALPHA) + new_val * ALPHA

    def play_episode(self, env):
        total_reward = 0.0
        state = env.reset()
        while True:
            _, action = self.best_value_and_action(state)
            new_state, reward, is_done, _ = env.step(action)
            total_reward += reward
            if is_done:
                break
            state = new_state
        return total_reward


if __name__ == "__main__":
    test_env = gym.make(ENV_NAME)
    agent = Agent()
    writer = SummaryWriter(comment="-q-learning")

    iter_no = 0
    best_reward = 0.0
    while True:
        iter_no += 1
        # 这种更新Q值的方法，关键在于这部，随机
        # 随机采样，遍历可能到达的状态，更新Q值
        # 与上一章的方法不同，上一章的方法每到到一个状态就会
        # 依据动作空间更新Q值，需要重复遍历一些不可能遍历执行的动作
        s, a, r, next_s = agent.sample_env()
        agent.value_update(s, a, r, next_s)

        reward = 0.0
        for _ in range(TEST_EPISODES):
            reward += agent.play_episode(test_env)
        reward /= TEST_EPISODES
        writer.add_scalar("reward", reward, iter_no)
        if reward > best_reward:
            print("Best reward updated %.3f -> %.3f" % (best_reward, reward))
            best_reward = reward
        if reward > 0.80:
            print("Solved in %d iterations!" % iter_no)
            break
    writer.close()
